from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import List

mcp = FastMCP("Cloud-init MCP Server")

"""
面向 cloud-init 软件的 MCP Server

所有 MCP tool 列表：
1. cloud_init_service_control - 控制 cloud-init 相关服务的启动、停止、重启和重载
   参数：service_name, action
   返回：统一的命令执行结果 JSON

2. cloud_init_service_status - 检查 cloud-init 相关服务的状态
   参数：service_name
   返回：统一的命令执行结果 JSON

3. cloud_init_oneshot_service_check - 检查 oneshot 类型服务的状态
   参数：service_name
   返回：统一的命令执行结果 JSON
"""

@mcp.tool()
def cloud_init_service_control(service_name: str, action: str) -> dict:
    """
    控制 cloud-init 相关服务的启动、停止、重启和重载操作
    
    参数:
        service_name: 服务名称，如 cloud-config.service, cloud-final.service, cloud-init.service 等
        action: 操作类型，支持 start|stop|restart|reload
        
    返回 JSON 结构:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    valid_actions = ['start', 'stop', 'restart', 'reload']
    if action not in valid_actions:
        return {
            "success": False,
            "command": f"systemctl {action} {service_name}",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Invalid action: {action}. Supported actions: {valid_actions}"
        }
    
    command = f"systemctl {action} {service_name}"
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command execution timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def cloud_init_service_status(service_name: str) -> dict:
    """
    检查 cloud-init 相关服务的状态
    
    参数:
        service_name: 服务名称，如 cloud-config.service, cloud-final.service, cloud-init.service 等
        
    返回 JSON 结构:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    command = f"systemctl status {service_name}"
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command execution timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

@mcp.tool()
def cloud_init_oneshot_service_check(service_name: str) -> dict:
    """
    检查 oneshot 类型服务的状态，特别处理 Type=oneshot 的服务
    
    参数:
        service_name: 服务名称，如 cloud-init-hotplugd.service
        
    返回 JSON 结构:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    # 首先检查服务类型
    check_type_command = f"systemctl show {service_name} -p Type --value"
    try:
        type_result = subprocess.run(check_type_command, shell=True, capture_output=True, text=True, timeout=30)
        service_type = type_result.stdout.strip()
        
        if service_type == "oneshot":
            # 对于 oneshot 服务，检查状态
            status_command = f"systemctl is-active {service_name}"
            status_result = subprocess.run(status_command, shell=True, capture_output=True, text=True, timeout=30)
            return {
                "success": status_result.returncode == 0,
                "command": f"{check_type_command} && {status_command}",
                "exit_code": status_result.returncode,
                "stdout": f"Service type: {service_type}\nActive status: {status_result.stdout}",
                "stderr": status_result.stderr
            }
        else:
            # 对于非 oneshot 服务，返回常规状态
            status_command = f"systemctl status {service_name}"
            status_result = subprocess.run(status_command, shell=True, capture_output=True, text=True, timeout=30)
            return {
                "success": status_result.returncode == 0,
                "command": f"{check_type_command} && {status_command}",
                "exit_code": status_result.returncode,
                "stdout": f"Service type: {service_type}\n{status_result.stdout}",
                "stderr": status_result.stderr
            }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": check_type_command,
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command execution timed out"
        }
    except Exception as e:
        return {
            "success": False,
            "command": check_type_command,
            "exit_code": -1,
            "stdout": "",
            "stderr": f"Command execution failed: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()